/*
 *    Copyright (c) 2014 Xiang Xu <xuxiang@mail.bnu.edu.cn>
 *
 *    Permission is hereby granted, free of charge, to any person
 *    obtaining a copy of this software and associated documentation
 *    files (the "Software"), to deal in the Software without
 *    restriction, including without limitation the rights to use,
 *    copy, modify, merge, publish, distribute, sublicense, and/or sell
 *    copies of the Software, and to permit persons to whom the
 *    Software is furnished to do so, subject to the following
 *    conditions:
 *
 *    The above copyright notice and this permission notice shall be
 *    included in all copies or substantial portions of the Software.
 *
 *    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
 *    EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
 *    OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
 *    NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
 *    HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
 *    WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 *    FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
 *    OTHER DEALINGS IN THE SOFTWARE.
 */

#include "arguments.h"
#include "pclwraper.h"

#include "facade_grammar.h"
#include "facade_model.h"

#include "svm.h"
#include "opencv2/opencv.hpp"
#include "boost/timer.hpp"
#include "boost/smart_ptr.hpp"

#include <iostream>
#include <vector>
#include <fstream>

#define Malloc(type,n) (type *)malloc((n)*sizeof(type))

using namespace std;

CvRect box;
bool drawing_box = false;
int n_terminals = 0;
int current_terminal = 0;
int facade_height;

bool ground_true = false;

bool learning = true;

CvScalar colors[] = {CV_RGB(255, 0, 0), CV_RGB(0, 255, 0), CV_RGB(0, 0, 255), CV_RGB(255, 255, 0),
                     CV_RGB(255, 0, 255), CV_RGB(0, 255, 255)
                    };

// query this after executing a function
Status status;

// organize variables
boost::shared_ptr<Arguments> arguments = boost::shared_ptr<Arguments>(new Arguments);
bool load_svm_model;
std::string svm_file;

// facade grammar and model
boost::shared_ptr<FacadeGrammar> facade_grammar;
boost::shared_ptr<FacadeModel> facade_model;

// svm variables
svm_problem problem;
svm_parameter parameter;
svm_model* model;
svm_node* x_space;
svm_node* y_space;

// trainning set
std::vector<std::vector<CvPoint> > training_samples;

// svm functions
void train();
void predict();
double grid_search(svm_problem* prob, svm_parameter* param, int nr_fold);

// A little subroutine to draw a box onto an image
void draw_box(IplImage* img, CvRect rect);

// Define our callback which we will install for mouse events
void my_mouse_callback(int event, int x, int y, int flags, void* param);

int main(int argc, char** argv)
{
    box = cvRect(-1, -1, 0, 0);

    // parse parameters
    status = arguments->Parse(argc, argv);

    if(!status.IsOk())
    {
        cout << status;
        exit(-1);
    }
    else
    {
        cout << "Argments Parsed success.\n";
    }

    // featch parameters
    load_svm_model = arguments->LoadQTable();
    svm_file = arguments->GetSVMModelFile();

    // Load the grammar
    facade_grammar = boost::shared_ptr<FacadeGrammar>(new FacadeGrammar);
    status = facade_grammar->ReadGrammar(arguments->GetGrammarFile());
    if (!status.IsOk())
	 {
        cout << status;
        exit(-1);
    }
    n_terminals = facade_grammar->get_terminal_symbol_size();
    training_samples.resize(n_terminals);

    // Initialize facade model and load the attributes
    facade_model = boost::shared_ptr<FacadeModel>(new FacadeModel(facade_grammar));
    facade_model->LoadGrids(arguments->GetGridFile());
    facade_height = facade_model->get_height();

    // Load the image
    std::ostringstream filename;
    filename << arguments->GetGridFile() << "curvate" << ".png";
    IplImage* img = cvLoadImage(filename.str().c_str());

    if(img == NULL)
    {
        std::cout << "Load image fail!\n";
        exit(-1);
    }

    // backup  the image
    IplImage* tmp = cvCloneImage(img);

    // Create a named window
    cvNamedWindow("SVM-Classifier");

    // Install the call back
    cvSetMouseCallback("SVM-Classifier", my_mouse_callback, (void*)img);

    // The main program loop.
    // Here we copy the working image to the 'temp' image,
    // and if the user is drawing, then put the currently
    // contemplated box onto that temp image. display the temp
    // image, and wait 15ms for a key stroke, then repeat ...
    while(1)
    {
        if(tmp != NULL)
            cvReleaseImage(&tmp);

        tmp = cvCloneImage(img);

        if(drawing_box)
            draw_box(tmp, box);

        cvShowImage("SVM-Classifier", tmp);

        int c = cvWaitKey(15);

        if((c & 255) == 27 || c == 'q')
        {
            std::cout << "Exiting ...\n";
            break;
        }

        switch((char)c)
        {
            case 'c':
                {
                    ++current_terminal;

                    if(current_terminal >= n_terminals)
                    {
                        current_terminal = 0;
                    }

                    std::cout << "current terminal: " << facade_grammar->get_terminal_symbol(current_terminal) << std::endl;
                }
                break;
            case 't':
                {
                    if(!load_svm_model)
                    {
                        boost::timer timer;
                        train();
                        std::cout << "Training time: " << timer.elapsed() << std::endl;
                    }
                    else
                    {
                        model = svm_load_model(svm_file.c_str());
                        std::cout << "SVM model loaded.\n";
                    }
                }
                break;
            case 'p':
                {
                    boost::timer timer;
                    predict();
                    std::cout << "Predict time: " << timer.elapsed() << std::endl;
                    learning = false;
                }
                break;
            case 'g':
                {
                    ground_true = true;
                }
                break;
            case 's':
                {
                    facade_model->RenderGrds(arguments->GetGridFile());
                    facade_model->SaveGrids(arguments->GetGridFile());
                    std::cout << "Classify results saved.\n";
                }
                break;
            default:
                break;
        }
    }

    // Be tidy
    cvReleaseImage(&img);
    cvReleaseImage(&tmp);
    cvDestroyWindow("SVM-Classifier");

    return 0;
}

// This is our mouse callback. If the user
// presses the left button, we start a box.
// when the under releases that button, then we
// add the box to he current image. When the mouse is
// dragged(width the button down) we resize the box
void my_mouse_callback(int event, int x, int y, int flags, void* param)
{
    IplImage* image = (IplImage*)param;

    switch(event)
    {
        case CV_EVENT_MOUSEMOVE:

            if(drawing_box)
            {
                box.width = x - box.x;
                box.height = y - box.y;
            }

            break;
        case CV_EVENT_LBUTTONDOWN:
            drawing_box = true;
            box = cvRect(x, y, 0, 0);
            break;
        case CV_EVENT_LBUTTONUP:
            drawing_box = false;

            if(box.width < 0)
            {
                box.x += box.width;
                box.width *= -1;
            }

            if(box.height < 0)
            {
                box.y += box.height;
                box.height *= -1;
            }

            draw_box(image, box);

            if(!ground_true)
            {
                std::cout << "Add " << box.width * box.height << " to terminal " << current_terminal << std::endl;

                for(int i = box.x; i < box.x + box.width; ++i)
                {
                    for(int j = box.y; j < box.y + box.height; ++j)
                    {
                        training_samples[current_terminal].push_back(cvPoint(i, facade_height - 1 - j));
                    }
                }

                std::cout << "terminal " << facade_grammar->get_terminal_symbol(current_terminal) << " has "
                          << training_samples[current_terminal].size() << " pixels.\n";
            }
            else
            {
                const int nr_class = svm_get_nr_class(model);

                for(int i = box.x; i < box.x + box.width; ++i)
                {
                    for(int j = box.y; j < box.y + box.height; ++j)
                    {
                        for(int k = 0; k < nr_class; ++k)
                        {
                            facade_model->SetProperty(i, j, k + FacadeFeature::FEATURE_NUMBER + 1, 0);
                        }

                        facade_model->SetProperty(i, facade_height - 1 - j, FacadeFeature::FEATURE_NUMBER + 1 + current_terminal, 1.0);
                        facade_model->SetProperty(i, facade_height - 1 - j, FacadeFeature::FEATURE_NUMBER, current_terminal);
                    }
                }
            }

            break;
        default:
            break;
    }
}

void draw_box(IplImage* img, CvRect rect)
{
    cvRectangle(img, cvPoint(box.x, box.y),
                cvPoint(box.x + box.width, box.y + box.height),
                colors[current_terminal],
                CV_FILLED
               );
}

void train()
{
    /* Construct SVM problem */
    std::cout << "Construct SVM problem ...\n";
    problem.l = 0;

    for(int i = 0; i < n_terminals; ++i)
    {
        problem.l += training_samples[i].size();
    }

    problem.y = Malloc(double, problem.l);
    problem.x = Malloc(struct svm_node*, problem.l);

    int elements = problem.l * (FacadeFeature::FEATURE_NUMBER + 1);
    x_space = Malloc(struct svm_node, elements);

    int sample_count = 0;
    int node_count = 0;

    for(int i = 0; i < n_terminals; ++i)
    {
        std::vector<CvPoint>& samples = training_samples[i];

        for(int j = 0; j < samples.size(); ++j)
        {
            // label
            problem.y[sample_count] = i;
            // features
            problem.x[sample_count++] = &x_space[node_count];

            int index = 0;

            for(int feature = 0; feature < FacadeFeature::FEATURE_NUMBER; ++feature)
            {
                x_space[node_count].value = facade_model->GetProperty(samples[j].x, samples[j].y, feature);
                x_space[node_count].index = ++index;
                ++node_count;
            }

            x_space[node_count++].index = -1;
        }
    }

    assert(sample_count == problem.l);

    /* Specify SVM parameters */
    std::cout << "Specify SVM parameters ...\n";
    // default values
    svm_parameter parameter;
    parameter.svm_type = C_SVC;
    parameter.kernel_type = RBF;
    parameter.degree = 3;
    parameter.gamma = 1;	// 1/num_features
    parameter.coef0 = 0;
    parameter.nu = 0.5;
    parameter.cache_size = 100;
    parameter.C = 1;
    parameter.eps = 1e-3;
    parameter.p = 0.1;
    parameter.shrinking = 0;
    parameter.probability = 1;
    parameter.nr_weight = 0;
    parameter.weight_label = NULL;
    parameter.weight = NULL;

    // check the parameters
    const char* error_infor = svm_check_parameter(&problem, &parameter);

    if(error_infor != NULL)
    {
        std::cout << error_infor << std::endl;
        exit(-1);
    }

    /* Cross Validate */
    std::cout << "Cross Validate ...\n";
    double accuracy = grid_search(&problem, &parameter, n_terminals);

    std::cout << "Training...\n";
    /* Training */
    model = svm_train(&problem, &parameter);

    svm_save_model(svm_file.c_str(), model);

    std::cout << "SVM model saved.\n";
    std::cout << "best c: " << parameter.C << " best g: " << parameter.gamma << " Accuracy: " << accuracy << std::endl;
}

void predict()
{
    std::cout << "Predict ...\n";
    int elements = facade_model->get_width() * facade_model->get_height() * (FacadeFeature::FEATURE_NUMBER + 1);
    y_space = Malloc(struct svm_node, elements);

    const int nr_class = svm_get_nr_class(model);
    double* probility_estimates = Malloc(double, nr_class);

    int sample_count = 0;

    for(int i = 0; i < facade_model->get_width(); ++i)
    {
        for(int j = 0; j < facade_model->get_height(); ++j)
        {
            int index = 0;

            for(int feature = 0; feature < FacadeFeature::FEATURE_NUMBER; ++feature)
            {
                y_space[sample_count].value = facade_model->GetProperty(i, j, feature);
                y_space[sample_count].index = ++index;
                ++sample_count;
            }

            y_space[sample_count++].index = -1;

            // double label = svm_predict(model, node);
            double label = svm_predict_probability(model, &y_space[sample_count - FacadeFeature::FEATURE_NUMBER - 1],
                                                   probility_estimates);

            facade_model->SetProperty(i, j, FacadeFeature::FEATURE_NUMBER, label);

            for(int k = 0; k < nr_class; ++k)
            {
                facade_model->SetProperty(i, j, k + FacadeFeature::FEATURE_NUMBER + 1, probility_estimates[k]);
            }
        }
    }

    svm_free_and_destroy_model(&model);
    svm_destroy_param(&parameter);
    free(problem.y);
    free(problem.x);
    free(x_space);
    free(y_space);
    free(probility_estimates);

    std::cout << "Done.\n";
}

double grid_search(svm_problem* prob, svm_parameter* param, int nr_fold)
{
    int log2c_begin = -5;
    int log2c_end = 15;
    int log2c_step = 2;
    int log2g_begin = 3;
    int log2g_end = -15;
    int log2g_step = -2;

    double* target = new double[prob->l];

    int best_correct = 0;
    double best_c = 1;
    double best_g = 1;
    double best_accuracy;

    for(int log2c = log2c_begin; log2c <= log2c_end; log2c += log2c_step)
    {
        for(int log2g = log2g_begin; log2g >= log2g_end; log2g += log2g_step)
        {
            param->C = pow(2.0, log2c);
            param->gamma = pow(2.0, log2g);

            svm_cross_validation(prob, param, nr_fold, target);
            int total_correct = 0;

            for(int i = 0; i  < prob->l; ++i)
                if(target[i] == prob->y[i])
                    ++total_correct;

            if(total_correct >= best_correct)
            {
                best_correct = total_correct;
                best_c = param->C;
                best_g = param->gamma;
            }

            best_accuracy = 100.0 * best_correct / prob->l;

            if(best_accuracy > 95.0)
                goto FINISH;
        }
    }

FINISH:
    param->C = best_c;
    param->gamma = best_g;

    delete[] target;
    best_accuracy = 100.0 * best_correct / prob->l;
    return best_accuracy;
}
